{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "OE4xjKWgtIX2"
   },
   "source": [
    "# Bunny ICP"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "L6p6VR5ktIX3"
   },
   "source": [
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the manipulation repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  \n",
    "\n",
    "More details are available [here](http://manipulation.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "KdXAerwV13rQ"
   },
   "outputs": [],
   "source": [
    "import importlib\n",
    "import os, sys\n",
    "from urllib.request import urlretrieve\n",
    "\n",
    "if 'google.colab' in sys.modules and importlib.util.find_spec('manipulation') is None:\n",
    "    urlretrieve(f\"http://manipulation.csail.mit.edu/scripts/setup/setup_manipulation_colab.py\",\n",
    "                \"setup_manipulation_colab.py\")\n",
    "    from setup_manipulation_colab import setup_manipulation\n",
    "    setup_manipulation(manipulation_sha='db19656799c11a58bc37dd20604ad38eafb09c62', drake_version='20200921', drake_build='nightly')\n",
    "\n",
    "# python libraries\n",
    "import numpy as np\n",
    "\n",
    "# Determine if this notebook is currently running as a notebook or a unit test.\n",
    "from IPython import get_ipython\n",
    "running_as_notebook = get_ipython() and hasattr(get_ipython(), 'kernel')\n",
    "\n",
    "# Start a single meshcat server instance to use for the remainder of this notebook.\n",
    "from meshcat.servers.zmqserver import start_zmq_server_as_subprocess\n",
    "server_args = []\n",
    "if 'google.colab' in sys.modules:\n",
    "    server_args = ['--ngrok_http_tunnel']\n",
    "proc, zmq_url, web_url = start_zmq_server_as_subprocess(server_args=server_args)\n",
    "\n",
    "import meshcat.geometry as g\n",
    "import meshcat.transformations as tf\n",
    "\n",
    "from pydrake.all import RigidTransform, RotationMatrix\n",
    "\n",
    "import open3d as o3d \n",
    "import meshcat\n",
    "\n",
    "from IPython.display import clear_output\n",
    "clear_output()\n",
    "\n",
    "from manipulation import FindResource"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "jigwRNW7tIYQ"
   },
   "source": [
    "## Problem Description\n",
    "In the lecture, we learned about the Iterative Closest Point (ICP) algorithm. In this exercise, you will implement the ICP algorithm to solve the standard Stanford Bunny problem!\n",
    "\n",
    "**These are the main steps of the exercise:**\n",
    "1. Implement the ```least_squares_transform``` function to optimize transformation given correspondence\n",
    "2. Implement the ```icp``` algorithm using the functions implemented above.\n",
    "\n",
    "Let's first visualize the point clouds of Stanford bunny in meshcat!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 51
    },
    "colab_type": "code",
    "id": "VnQvdI6nOl4d",
    "outputId": "866d6e69-d188-4669-8b01-825d8e616b0d"
   },
   "outputs": [],
   "source": [
    "# Visualize Stanford Bunny \n",
    "pcd = o3d.io.read_point_cloud(FindResource(\"models/bunny/bun_zipper_res2.ply\"))\n",
    "pointcloud_model = np.asarray(pcd.points).transpose()\n",
    "\n",
    "# First, clean the origin a bit to define nominal pose. \n",
    "X = np.array([[1., 0., 0., 0.0],\n",
    "              [0., np.cos(np.pi/2), -np.sin(np.pi/2), 0.],\n",
    "              [0., np.sin(np.pi/2), np.cos(np.pi/2), -0.05]])\n",
    "Xtemp = RigidTransform(X)\n",
    "X = np.array([[np.cos(np.pi/2), -np.sin(np.pi/2), 0, 0.],\n",
    "              [np.sin(np.pi/2), np.cos(np.pi/2), 0., 0.],\n",
    "              [0., 0., 1., 0.]])\n",
    "X = RigidTransform(X).multiply(Xtemp)\n",
    "\n",
    "pointcloud_model = X.multiply(pointcloud_model).T\n",
    "\n",
    "# Next, set up a a distrubance \n",
    "X = np.array([[1., 0., 0., -1e-1],\n",
    "              [0., np.cos(np.pi/6), -np.sin(np.pi/6), 1e-1],\n",
    "              [0., np.sin(np.pi/6), np.cos(np.pi/6), 1e-1]])\n",
    "X = RigidTransform(X)\n",
    "pointcloud_scene = X.multiply(pointcloud_model.T).T\n",
    "\n",
    "def generate_color_mat(color_vec, shape):\n",
    "  color_mat = np.tile(np.array(color_vec).astype(np.float32).reshape(3,1), (1, shape[1]))\n",
    "  return color_mat\n",
    "\n",
    "vis = meshcat.Visualizer(zmq_url)\n",
    "vis.delete()\n",
    "vis[\"/Background\"].set_property('visible', False)\n",
    "#vis[\"/Cameras/default/\"].set_transform(tf.translation_matrix([0, 0, 1]))\n",
    "vis[\"/Cameras/default/rotated/<object>\"].set_property(\"zoom\", 10.5)\n",
    "\n",
    "vis[\"red_bunny\"].set_object(g.PointCloud(pointcloud_model.T, generate_color_mat([1, 0, 0], pointcloud_model.T.shape), size=0.01))\n",
    "vis[\"blue_bunny\"].set_object(g.PointCloud(pointcloud_scene.T, generate_color_mat([0, 0, 1], pointcloud_scene.T.shape), size=0.01))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "abo92_2stIYW"
   },
   "source": [
    "## Point cloud registration with known correspondences\n",
    "\n",
    "In this section, you will follow the [derivation](http://manipulation.csail.mit.edu/pose.html#section3) to solve the optimization problem below. \n",
    "\n",
    "\\begin{align}\n",
    "    \\min_{p\\in\\mathbb{R}^3,R\\in\\mathbb{R}^{3\\times3}} \\quad & \\sum_{i=1}^{N_s} \\| p + R\n",
    "    \\hspace{.1cm} {^Op^{m_{c_i}}} - p^{s_i}\\|^2, \\\\ s.t. \\quad & RR^T =\n",
    "    I, \\quad \\det(R)=1\\end{align}\n",
    "    \n",
    "The goal is to find the transform that registers the point clouds of the model and the scene, assuming the correspondence is known.  You may refer to the implementation from [colab](https://colab.research.google.com/github/RussTedrake/manipulation/blob/master/pose.ipynb#scrollTo=AHfxMwrvb1mz) and the explanation from [textbook](http://manipulation.csail.mit.edu/pose.html#section4).\n",
    "\n",
    "In the cell below, implement the ```least_squares_transform``` nethod."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ll_FlqVotIYX"
   },
   "outputs": [],
   "source": [
    "def least_squares_transform(scene, model) -> RigidTransform:\n",
    "    '''\n",
    "    Calculates the least-squares best-fit transform that maps corresponding\n",
    "    points scene to model.\n",
    "    Args:\n",
    "      scene: Nx3 numpy array of corresponding points\n",
    "      model: Nx3 numpy array of corresponding points\n",
    "    Returns:\n",
    "      X_BA: A RigidTransform object that maps point_cloud_A on to point_cloud_B \n",
    "            such that\n",
    "                        X_BA.multiply(model) ~= scene,\n",
    "    '''\n",
    "    X_BA = RigidTransform()\n",
    "    ##################\n",
    "    # your code here\n",
    "    ##################\n",
    "    return X_BA"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "IejlqJ3vtIYg"
   },
   "source": [
    "## Point correspondence from closest point\n",
    "The ```least_squares_transform``` function assumes that the point correspondence is known. Unfortunately, this is often not the case, so we will have to estimate the point correspondence as well. A common heuristics for estimating point correspondence is the closest point/nearest neighbor. \n",
    "\n",
    "We have implemented the closest neighbors using [Open3d's implementation](http://www.open3d.org/docs/release/python_api/open3d.geometry.KDTreeFlann.html), which uses [k-d trees](https://en.wikipedia.org/wiki/K-d_tree)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "_-bGj1a1OkbU"
   },
   "outputs": [],
   "source": [
    "def nearest_neighbors(scene, model):\n",
    "    '''\n",
    "    Find the nearest (Euclidean) neighbor in model for each\n",
    "    point in scene\n",
    "    Args:\n",
    "        scene: Nx3 numpy array of points\n",
    "        model: Mx3 numpy array of points\n",
    "    Returns:\n",
    "        distances: (N, ) numpy array of Euclidean distances from each point in\n",
    "            scene to its nearest neighbor in model.\n",
    "        indices: (N, ) numpy array of the indices in model of each\n",
    "            scene point's nearest neighbor - these are the c_i's\n",
    "    '''\n",
    "    distances = np.empty(scene.shape[0], dtype=float)\n",
    "    indices = np.empty(scene.shape[0], dtype=int)\n",
    "    \n",
    "    kdtree = o3d.geometry.KDTreeFlann(model.T)\n",
    "    for i in range(model.shape[0]):\n",
    "        nn = kdtree.search_knn_vector_3d(scene[i,], 1)\n",
    "        indices[i] = nn[1][0]\n",
    "        distances[i] = np.linalg.norm(scene[i,:] - model[indices[i],:])\n",
    "\n",
    "    return distances, indices"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KtvN0kBntIYo"
   },
   "source": [
    "## Iterative Closest Point (ICP)\n",
    "Now you should be able to register two point clouds iteratively by first finding/updating the estimate of point correspondence with ```nearest_neighbors``` and then computing the transform using ```least_squares_transform```. You may refer to the explanation from [textbook](http://manipulation.csail.mit.edu/pose.html#section4).\n",
    "\n",
    "**In the cell below, complete the implementation of ICP algorithm using the  ```nearest_neighbors``` and ```least_squares_transform``` methods from above.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "wETDMjk4tIYp"
   },
   "outputs": [],
   "source": [
    "def icp(scene, model, max_iterations=20, tolerance=1e-3, anim=None):\n",
    "    '''\n",
    "    Perform ICP to return the correct relative transform between two set of points.\n",
    "    Args:\n",
    "        scene: Nx3 numpy array of points\n",
    "        model: Mx3 numpy array of points\n",
    "        max_iterations: max amount of iterations the algorithm can perform.\n",
    "        tolerance: tolerance before the algorithm converges.\n",
    "    Returns:\n",
    "      X_BA: A RigidTransform object that maps point_cloud_A on to point_cloud_B \n",
    "            such that\n",
    "                        X_BA.multiply(model) ~= scene,\n",
    "      mean_error: Mean of all pairwise distances. \n",
    "      num_iters: Number of iterations it took the ICP to converge. \n",
    "    '''\n",
    "    X_BA = RigidTransform()\n",
    "\n",
    "    mean_error = 0\n",
    "    num_iters = 0\n",
    "    prev_error = 0\n",
    "    \n",
    "    while True:\n",
    "        num_iters += 1  \n",
    "          \n",
    "        # your code here\n",
    "        ##################\n",
    "\n",
    "        mean_error = np.inf # Modify to add mean error.\n",
    "        ##################\n",
    "\n",
    "        if abs(mean_error - prev_error) < tolerance or num_iters >= max_iterations:\n",
    "            break\n",
    "\n",
    "        prev_error = mean_error\n",
    "\n",
    "        if anim is not None:\n",
    "          with anim.at_frame(vis, 5 * num_iters) as frame:\n",
    "            frame[\"red_bunny\"].set_transform(X_BA.GetAsMatrix4())\n",
    "\n",
    "    return X_BA, mean_error, num_iters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "WChfoIVWtIYy"
   },
   "source": [
    "Now you should be able to visualize the registration of the Stanford bunny! Have fun!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "P8oIuMsCDMJM"
   },
   "outputs": [],
   "source": [
    "from meshcat.animation import Animation\n",
    "\n",
    "anim = Animation()\n",
    "icp(pointcloud_scene, pointcloud_model, max_iterations=30, tolerance=1e-5, anim=anim)\n",
    "\n",
    "vis.set_animation(anim)\n",
    "vis[\"/Animations/default\"].set_property(\"timeScale\", 10.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ucRnypactIY2"
   },
   "source": [
    "## How will this notebook be Graded?##\n",
    "\n",
    "If you are enrolled in the class, this notebook will be graded using [Gradescope](www.gradescope.com). You should have gotten the enrollement code on our announcement in Piazza. \n",
    "\n",
    "For submission of this assignment, you must do two things. \n",
    "- Download and submit the notebook `bunny_icp.ipynb` to Gradescope's notebook submission section, along with your notebook for the other problems.\n",
    "- Copy and Paste your answer to the kinematic singularity problem to Gradescope's written submission section. \n",
    "\n",
    "We will evaluate the local functions in the notebook to see if the function behaves as we have expected. For this exercise, the rubric is as follows:\n",
    "- [3 pts] `least_squares_transform` must be implemented correctly. \n",
    "- [3 pts] `icp` must be implemented correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "lfmnYSMItIY3",
    "outputId": "21d3a0f9-0e63-4eb7-cc7f-d95677e38c5e"
   },
   "outputs": [],
   "source": [
    "from manipulation.exercises.pose.test_icp import TestICP\n",
    "from manipulation.exercises.grader import Grader \n",
    "\n",
    "Grader.grade_output([TestICP], [locals()], 'results.json')\n",
    "Grader.print_test_results('results.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "colab": {
   "collapsed_sections": [],
   "name": "bunny_icp.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}